﻿import os
import glob
import random
import shutil
from PIL import Image

# 对所有图片进行RGB转化，并且统一调整到一致大小，但不让图片发生变形或扭曲，划分了训练集和测试集"FromPILimportImage
if __name__ == "__main__":
    test_split_ratio = 0.05
    desired_size = 128  # 图片缩放后的统一大小
    raw_path = "./raw"

    current_working_directory = os.getcwd()
    file_dir_path = os.path.join(
        current_working_directory, "Pytorch", "ResNet", "Fruit"
    )

    dirs = glob.glob(os.path.join(file_dir_path, "raw", "*"))
    dirs = [d for d in dirs if os.path.isdir(d)]
    print(f"Totally {len(dirs)} classes:{dirs}")
    for full_path in dirs:
        # 对每个类别单独处理
        path = full_path.split("\\")[-1]
        os.makedirs(file_dir_path + f"\\train\\{path}", exist_ok=True)
        os.makedirs(file_dir_path + f"\\test\\{path}", exist_ok=True)
        files = glob.glob(os.path.join(full_path, "*.jpg"))
        # files += glob.glob(os.path.join(full_path, "*.JPG"))
        files += glob.glob(os.path.join(full_path, "*.png"))
        random.shuffle(files)

        boundary = int(len(files) * test_split_ratio)  # 训练集和测试集的边界
        for i, file in enumerate(files):
            img = Image.open(file).convert("RGB")
            old_size = img.size  # old_size[0] is in (width,height） format
            ratio = float(desired_size) / max(old_size)
            new_size = tuple([int(x * ratio) for x in old_size])
            im = img.resize(new_size, Image.Resampling.LANCZOS)

            new_im = Image.new("RGB", (desired_size, desired_size))
            new_im.paste(
                im,
                ((desired_size - new_size[0]) // 2, (desired_size - new_size[1]) // 2),
            )
            assert new_im.mode == "RGB"
            if i <= boundary:
                new_im.save(
                    os.path.join(
                        file_dir_path,
                        "test",
                        path,
                        file.split("\\")[-1].split(".")[0] + ".jpg",
                    )
                )
            else:
                new_im.save(
                    os.path.join(
                        file_dir_path,
                        "train",
                        path,
                        file.split("\\")[-1].split(".")[0] + ".jpg",
                    )
                )
    test_files = glob.glob(os.path.join("test", "*", "*.jpg"))
    train_files = glob.glob(os.path.join("train", "*", "*.jpg"))